import numpy as np

class Adam:
    def __init__(self, shape: tuple):
        self.beta1 = 0.9
        self.beta2 = 0.999
        self.epsilon = 1e-8
        self.timestep = 0
        self.m = np.zeros(shape)
        self.v = np.zeros(shape)
    
    def step(self, params: np.ndarray, grads: np.ndarray, lr: float):
        self.timestep += 1
        self.m = self.m * self.beta1 + grads * (1.0 - self.beta1)
        self.v = self.v * self.beta2 + (grads ** 2) * (1.0 - self.beta2)
        
        m_hat = self.m / (1.0 - self.beta1 ** self.timestep)
        v_hat = self.v / (1.0 - self.beta2 ** self.timestep)
        
        update = m_hat / (np.sqrt(v_hat) + self.epsilon)
        
        params -= update * lr